from selenium import webdriver
from selenium.webdriver import ActionChains
import os
import base64
from PIL import Image
import time
from tqdm import tqdm
# import cv2
from reportlab.lib.pagesizes import A4, portrait, landscape
from reportlab.pdfgen import canvas
from PyPDF2 import PdfFileReader as reader, PdfFileWriter as writer
import sys
import re
# reload(sys)
# sys.setdefaultencoding('utf-8')


class CTGU(object):
    def __init__(self,
                 username="你的学号",
                 password="你的密码"):
        """
        登录图书馆
        """
        # --- 配置chromedriver ---
        driver = webdriver.Chrome()
        # driver = webdriver.Chrome(executable_path=r"C:\Program Files\Google\Chrome\Application\chromedriver.exe")
        # --- 图书馆 ---
        login_page = 'http://cas.lib.ctgu.edu.cn/cas/login'
        driver.get(login_page)
        # --- 输入学号 ---
        username_input = driver.find_element_by_css_selector(
            '#username[name="username"]')
        username_input.clear()
        username_input.send_keys(username)
        # --- 输入密码 ---
        password_input = driver.find_element_by_css_selector(
            '#password[name="password"]')
        password_input.clear()
        password_input.send_keys(password)
        # --- 滑动验证 ---
        slide_btn = driver.find_element_by_id('btn')
        action = ActionChains(driver)
        action.click_and_hold(slide_btn)
        action.move_by_offset(290-54, 0)
        action.release()
        action.perform()
        # --- 点击登录 ---
        login_btn = driver.find_element_by_css_selector('.fm1_submit')
        action = ActionChains(driver)
        action.click(login_btn)
        action.perform()
        self.driver = driver

    def img2bs4(self):
        """
        图片链接转 base64
        """
        js = """
                _fetch = function(i,src){
                return fetch(src).then(function(response) {
                    if(!response.ok) throw new Error("No image in the response");
                    var headers = response.headers;
                    var ct = headers.get('Content-Type');
                    var contentType = 'image/png';
                    if(ct !== null){
                    contentType = ct.split(';')[0];
                    }
                    
                    return response.blob().then(function(blob){
                    return {
                        'blob': blob,
                        'mime': contentType,
                        'i':i,
                    };
                    });
                });
                };
                
                _read = function(response){
                return new Promise(function(resolve, reject){
                    var blob = new Blob([response.blob], {type : response.mime});
                    var reader = new FileReader();
                    reader.onload = function(e){
                    resolve({'data':e.target.result, 'i':response.i});
                    };
                    reader.onerror = reject;
                    reader.readAsDataURL(blob);
                });
                };
                
                _replace = function(){
                    for (var i = 0, len = q.length; i < len; i++) {imgs[q[i].item].src = q[i].data;}
                }
                
                var q = [];
                var imgs = document.querySelectorAll('img');
                for (var i = 0, len = imgs.length; i < len; i++) {
                        _fetch(i,imgs[i].src).then(_read).then(function(data){
                    q.push({
                    'data': data.data,
                    'item': data.i,
                    });
                });
                    }
                setTimeout(_replace, 1000 );
                """
        self.driver.execute_script(js)

    def nextPage(self):
        """
        自动翻页
        """
        next_btn = self.driver.find_element_by_id('aNextPage')
        action = ActionChains(self.driver)
        action.click(next_btn)
        action.perform()

    def setPage(self, index):
        """
        转到某一页
        """
        js = "SetPage({})".format(index)
        self.driver.execute_script(js)

    def downloadImg(self, index, save_path, sleep_seconds):
        """
        下载图片
        """
        self.img2bs4()
        time.sleep(sleep_seconds)
        img_bs64 = self.driver.find_element_by_id(
            "Img{}".format(index+1)).get_attribute('src')
        img_data = base64.urlsafe_b64decode(
            img_bs64[21:-1] + '=' * (4 - len(img_bs64) % 4))
        save_name = os.path.join(save_path, '{:0=3}.png'.format(index+1))
        with open(save_name, 'wb') as f:
            f.write(img_data)
        f.close()

    def savePDF(self, img_path, pdf_path):
        """
        生成pdf
        """
        pages = 0
        (w, h) = portrait(A4)
        c = canvas.Canvas(pdf_path, pagesize=portrait(A4))
        l = os.listdir(img_path)
        l.sort(key=lambda x: int(x[:-4]))
        for i in l:
            f = img_path + os.sep + str(i)
            c.drawImage(f, 0, 0, w, h)
            c.showPage()
            pages = pages + 1
        c.save()

    def makeCategory(self, pdf_path):
        """
        生成目录
        """
        book = reader(pdf_path)
        # 创建PdfFileWriter对象，并用拷贝reader对象进行初始化
        pdf = writer()
        pdf.cloneDocumentFromReader(book)
        cat_div = self.driver.find_elements("css selector", '.cat_con a')
        for div in cat_div:
            catText = div.get_attribute("innerHTML").strip(' ')
            carPage = int(re.findall(
                r'"([^"]*)"', div.get_attribute("onclick"))[0])
            pdf.addBookmark(catText, carPage-1)

        # 保存修改后的PDF文件内容到文件中
        with open(pdf_path, 'wb') as fout:
            pdf.write(fout)

    def download(self,
                 paperUrl,  # 论文阅读地址
                 category,  # 是否生成目录
                 save_path, # 保存路径
                 sleep_seconds # 睡眠时间，防止图片因网速过慢或操作过快而加载失败
                 ):
        # --- 转到阅读界面 ---
        self.driver.get(paperUrl)
        time.sleep(2)
        self.driver.get(paperUrl)
        # --- 定位iframe ---
        iframe = self.driver.find_element_by_id('frmReader')
        self.driver.switch_to.frame(iframe)
        # --- 获取总页数 ---
        totalPageNum = int(self.driver.find_element_by_id(
            'hidEnd').get_attribute('value'))
        img_save_path = os.path.join(save_path, 'src')
        if not os.path.exists(img_save_path):
            os.makedirs(img_save_path)
        # --- 下载图片 ---
        for index in tqdm(range(totalPageNum)):
            # self.nextPage()
            if index % 10 == 0:
                time.sleep(15)
            self.nextPage()
            self.img2bs4()
            self.downloadImg(index, img_save_path, sleep_seconds)

        pdf_save_path = os.path.join(save_path, 'output.pdf')
        self.savePDF(img_save_path, pdf_save_path)
        if category:
            self.makeCategory(pdf_save_path)
        print('论文下载完毕！')
